package com.xz.xzmybatisagent.utils;

import java.io.File;
import java.net.MalformedURLException;
import java.net.URL;
import java.net.URLClassLoader;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.*;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;

/**
 * <p>
 * 扫描指定包下的所有类。
 *  TODO 排除目录
 * </p>
 *
 * @author DONG.WANG17@HAND-CHINA.COM
 * @since 2025/3/24-03-24 14:12
 */
public class ClassScanner {
    private static final String DEFAULT_FILTER_PREFIX = "java.main.";

    /**
     * 获取包下所有类（指定是否遍历子包）
     *
     * @param packageName 包名
     * @return 所有类名（获取到的全限定名称是从jar包或 build out 目录下的 classes 目录下开始的）
     */
    public static List<String> getClassName(String packageName) {
        return getClassName(packageName, true, Collections.emptyList(), true);
    }

    /**
     * 获取包下所有类（指定是否遍历子包，指定是否排除某些类）
     *
     * @param packageName  包名
     * @param childPackage 是否遍历子包
     * @return 所有类名（获取到的全限定名称是从jar包或 build out 目录下的 classes 目录下开始的）
     */
    public static List<String> getClassName(String packageName, boolean childPackage) {
        return getClassName(packageName, childPackage, Collections.emptyList(), true);
    }

    /**
     * 获取包下所有类（指定是否遍历子包，指定是否排除某些类，指定是否包含内部类）
     *
     * @param packageName         包名
     * @param childPackage        是否遍历子包
     * @param excludedClasses     排除列表
     * @param includeInnerClasses 是否包含内部类
     * @return 所有类名（获取到的全限定名称是从jar包或 build out 目录下的 classes 目录下开始的）
     */
    public static List<String> getClassName(String packageName, boolean childPackage, List<String> excludedClasses, boolean includeInnerClasses) {
        return getClassName(packageName, childPackage, excludedClasses, includeInnerClasses, DEFAULT_FILTER_PREFIX);
    }

    /**
     * 获取包下所有类（指定是否遍历子包，指定是否排除某些类，指定是否包含内部类，指定过滤前缀）
     *
     * @param packageName         包名
     * @param childPackage        是否遍历子包
     * @param excludedClasses     排除列表
     * @param includeInnerClasses 是否包含内部类
     * @param filterPrefix        过滤前缀（从 classes 目录开始的前缀。通过 package.package.package 的形式指定）
     * @return 所有类名（获取到的全限定名称是从jar包或 build out 目录下的 classes 目录下开始的）
     */
    public static List<String> getClassName(String packageName, boolean childPackage, List<String> excludedClasses, boolean includeInnerClasses, String filterPrefix) {
        List<String> fileNames = new ArrayList<>();
        //ClassLoader loader = Thread.currentThread().getContextClassLoader();
        String packagePath = packageName.replace(".", "/");
        //URL url = loader.getResource(packagePath);
        final URL url = ClassScanner.class.getProtectionDomain().getCodeSource().getLocation();

        Set<String> excludedSet = new HashSet<>(excludedClasses != null ? excludedClasses : Collections.emptyList());

        try {
            if (url != null) {
                String type = url.getProtocol();
                if ("file".equals(type)) {
                    fileNames = getClassNameByFile(url.getPath(), childPackage, excludedSet, includeInnerClasses, filterPrefix);
                } else if ("jar".equals(type)) {
                    fileNames = getClassNameByJar(url.getPath(), childPackage, excludedSet, includeInnerClasses, filterPrefix);
                }
            } else {
                fileNames = getClassNameByJars(((URLClassLoader) ClassScanner.class.getClassLoader()).getURLs(), packagePath, childPackage, excludedSet, includeInnerClasses, filterPrefix);
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
        return fileNames;
    }

    private static List<String> getClassNameByFile(String filePath, boolean childPackage, Set<String> excludedSet, boolean includeInner, String filterPrefix) {
        List<String> classNames = new ArrayList<>();
        File file = new File(filePath);
        File[] files = file.listFiles();
        if (files == null) return classNames;

        for (File child : files) {
            if (child.isDirectory() && childPackage) {
                classNames.addAll(getClassNameByFile(child.getAbsolutePath(), true, excludedSet, includeInner, filterPrefix));
            } else if (child.getName().endsWith(".class")) {
                String normalizedPath = child.getAbsolutePath().replace(File.separatorChar, '/');

                int classIdx = normalizedPath.indexOf("/classes/");

                if (classIdx == -1) continue;

                String className = normalizedPath.substring(classIdx + 9)
                        .replace(".class", "")
                        .replace("/", ".");

                // 根据前缀去除
                if (className.startsWith(formatPrefix(filterPrefix))) {
                    className = className.substring(formatPrefix(filterPrefix).length());
                }

                if (!includeInner && className.contains("$")) continue;

                if (excludedSet.contains(className)) continue;

                classNames.add(className);
            }
        }
        return classNames;
    }

    private static List<String> getClassNameByJar(String jarPath, boolean childPackage, Set<String> excludedSet, boolean includeInner, String filterPrefix) {
        List<String> classNames = new ArrayList<>();
        String[] jarInfo = jarPath.split("!");
        String jarFile = jarInfo[0].substring(jarInfo[0].indexOf("/"));
        String packagePath = jarInfo[1].substring(1).replace("/", ".");

        try (JarFile jar = new JarFile(jarFile)) {
            Enumeration<JarEntry> entries = jar.entries();
            while (entries.hasMoreElements()) {
                JarEntry entry = entries.nextElement();
                String entryName = entry.getName().replace("/", ".");
                if (!entryName.endsWith(".class")) continue;

                String className = entryName.substring(0, entryName.length() - 6);

                // 根据前缀去除
                if (className.startsWith(formatPrefix(filterPrefix))) {
                    className = className.substring(formatPrefix(filterPrefix).length());
                }

                String currentPackage = className.contains(".") ? className.substring(0, className.lastIndexOf(".")) : "";

                boolean packageMatch = childPackage ? className.startsWith(packagePath) : currentPackage.equals(packagePath);

                if (packageMatch) {
                    if (!includeInner && className.contains("$")) continue;
                    if (excludedSet.contains(className)) continue;
                    classNames.add(className);
                }
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
        return classNames;
    }

    private static List<String> getClassNameByJars(URL[] urls, String packagePath, boolean childPackage, Set<String> excludedSet, boolean includeInner, String filterPrefix) {
        List<String> classNames = new ArrayList<>();
        for (URL url : urls) {
            String path = url.getPath();
            if (path.endsWith("classes/")) continue;
            String jarPath = path + "!/" + packagePath;
            classNames.addAll(getClassNameByJar(jarPath, childPackage, excludedSet, includeInner, filterPrefix));
        }
        return classNames;
    }

    private static String formatPrefix(String prefix) {
        if (prefix == null) return "";
        return prefix.endsWith(".") ? prefix : prefix + ".";
    }

    public static List<String> getClassNameFromJar(String jarPath, ClassScanConfig scanConfig) {
        List<String> classNames = new ArrayList<>();
        try (JarFile jar = new JarFile(jarPath)) {
            Enumeration<JarEntry> entries = jar.entries();
            while (entries.hasMoreElements()) {
                JarEntry entry = entries.nextElement();
                if (!entry.getName().endsWith(".class")) continue;
                // 获取类名
                String className = entry.getName().substring(0, entry.getName().length() - ".class".length()).replace("/", ".");
                if (className.startsWith(scanConfig.getRemovePrefix())) {
                    className = className.substring(scanConfig.getRemovePrefix().length());
                }

                if (scanConfig.getPackagePrefix() != null && !scanConfig.getPackagePrefix().trim().isEmpty()) {
                    if (!className.startsWith(scanConfig.getPackagePrefix())) continue;
                    if (!scanConfig.getChildPackage() && !className.substring(0, className.lastIndexOf(".")).equals(scanConfig.getPackagePrefix()))
                        continue;
                }
                if (className.startsWith(scanConfig.getFilterPrefix())) continue;
                if (!scanConfig.getIncludeInnerClasses() && className.contains("$")) continue;
                if (scanConfig.getExcludedClasses().contains(className)) continue;
                if (scanConfig.getExcludedDirectories().stream().anyMatch((s) -> entry.getName().contains(s))) continue;
                classNames.add(className);
            }
        } catch (Exception ignored) {
            //
        }
        return classNames;
    }

    public static List<String> getClassNameFromJar(URL jarUrl, ClassScanConfig scanConfig) {
        return getClassNameFromJar(jarUrl.getPath(), scanConfig);
    }

    public static List<String> getClassNameFromJar(Path path, ClassScanConfig scanConfig) throws MalformedURLException {
        return getClassNameFromJar(path.toUri().toURL(), scanConfig);
    }

    public static void main(String[] args) {
        String jarPath = "C:\\codespace\\IDEA\\plugin\\xz-utils\\src\\agent\\xz-mybatis-agent.jar";
        ClassScanConfig scanConfig = new ClassScanConfig();
        scanConfig.setPackagePrefix("com.xz.xzmybatisagent");
        scanConfig.setExcludedClasses(Arrays.asList("com.xz.xzmybatisagent.AgentmainAgent", "com.xz.xzmybatisagent.InterceptorClassFileTransformer", "com.xz.xzmybatisagent.SqlSender", "com.xz.xzmybatisagent.SqlReceiver", "com.xz.xzmybatisagent.SqlVo"));
        scanConfig.setChildPackage(true);
        scanConfig.setExcludedDirectories(Collections.singletonList("com/xz/xzmybatisagent/utils"));
        scanConfig.setIncludeInnerClasses(false);
        final List<String> classNameFromJar = getClassNameFromJar(jarPath, scanConfig);
        classNameFromJar.forEach(System.out::println);
    }
}
